import os
import torch.nn as nn
from PIL import Image
from torchvision import transforms
from matplotlib import pyplot as plt
from common_tools import transform_invert, set_seed

set_seed(5)

# ===================================================== load images =================================
# os.path.abspath(__file__) - 这行代码获取当前执行的Python文件的绝对路径。__file__是一个内置变量，表示当前文件的路径。os.path.abspath()函数确保返回的路径是绝对路径。
# os.path.dirname(os.path.abspath(__file__)) - 这行代码获取上述绝对路径的目录部分。也就是说，它将返回当前Python文件所在的文件夹的路径。
# os.path.join(os.path.dirname(os.path.abspath(__file__)), 'imgs', 'lena.png') - 这行代码使用os.path.join()函数将目录路径和两个子路径('imgs'和'lena.png')合并为一个完整的文件路径。假设当前Python文件位于/home/user/project/script.py，那么这行代码将返回/home/user/project/imgs/lena.png
path_img = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'imgs', 'lena.png')
print(path_img)
img = Image.open(path_img).convert('RGB')

# convert to tensor
img_transform = transforms.Compose([transforms.ToTensor()])
img_tensor = img_transform(img)
# 添加batch维度
img_tensor.unsqueeze_(dim=0) #C*H*W to B*C*H*W
# squ
# print("suqueeze前尺寸：{} \n, suqueeze后尺寸:{}".format(img_tensor.unsqueeze(dim=0).shape, img_tensor.squeeze().shape))

# ============================= create convolution layer ===========================================

# flag = 1
flag = 0
if flag:
    conv_layer = nn.Conv2d(3, 1, 3)
    #初始化卷积层权重 https://zhuanlan.zhihu.com/p/458373836
    nn.init.xavier_normal_(conv_layer.weight.data)

    # calculation
    img_conv = conv_layer(img_tensor)

# ======================== transposed layer =========================================
# flag = 0
if not flag:
    conv_layer = nn.ConvTranspose2d(3, 1, 3, stride=2) # input:(input_channel, output_channel, size)
    # 初始化网络层权权值
    nn.init.xavier_normal_(conv_layer.weight.data)

    # calculation
    img_conv = conv_layer(img_tensor)

# ======================== visualization =========================================
print("卷积前尺寸：{} \n, 卷积后尺寸:{}".format(img_tensor.shape, img_conv.shape))
img_conv = transform_invert(img_conv[0, 0:1, ...], img_transform)
img_raw = transform_invert(img_tensor.squeeze(), img_transform)
plt.subplot(122).imshow(img_conv, cmap='gray')
plt.subplot(121).imshow(img_raw)
plt.show()

